import os
from torchvision import transforms,datasets
from PIL import Image
import torch
from torch.utils import data
import random
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim



cat_path=r'C:\Users\27437\Desktop\Python\minist\train\Cat'
dog_path=r'C:\Users\27437\Desktop\Python\minist\train\Dog'
cat_list=os.listdir(cat_path)
dog_list=os.listdir(dog_path)

# o=cv2.imread(cat_path+'\\'+cat_list[2])  #查看图片大小和通道数，shape加上（）出来的元素就变成元组了
# o=torch.tensor(o)
# print(o)
# print(o.size())
print(len(cat_list))  #打印样本容量
print(len(dog_list))  #打印样本容量


test_rate=0.2
cat_num=int(test_rate*len(cat_list))
dog_num=int(test_rate*len(dog_list))

cat_test=r'C:\Users\27437\Desktop\Python\minist\test\cat'
dog_test=r'C:\Users\27437\Desktop\Python\minist\test\dog'

def random_cat_move_files():
    for i in range(0,cat_num+1):
    # 随机选择一个文件
        random_cat_file = random.choice(cat_list)
        cat_list.remove(random_cat_file)
        source_cat_path = os.path.join(cat_path, random_cat_file)

        destination_cat_path = os.path.join(cat_test, random_cat_file)
    # 移动文件
        os.rename(source_cat_path, destination_cat_path)

def random_dog_move_files():
    for i in range(0,dog_num+1):
        random_dog_file = random.choice(dog_list)
        dog_list.remove(random_dog_file)
        source_dog_path=os.path.join(dog_path,random_dog_file)
        destination_dog_path=os.path.join(dog_test, random_dog_file)
        os.rename(source_dog_path, destination_dog_path)


# random_cat_move_files()
# random_dog_move_files()


#修改文件名
# cat_dog_list=os.listdir(a)
# for i ,o in enumerate (cat_dog_list):
#     file_path = os.path.join(a, o)
#     new='dog_'+str(i)+'.jpg'
#     os.rename(file_path, os.path.join(directory, new))


transforms = transforms.Compose(
[

transforms.RandomResizedCrop(150),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
                            std=[0.229, 0.224, 0.225])
]

)
train=r'C:\Users\27437\Desktop\Python\minist\train'
test=r'C:\Users\27437\Desktop\Python\minist\test'
train_data = datasets.ImageFolder(train, transforms)
test_data=datasets.ImageFolder(test, transforms)
#
# print(train_data[9999][0])  #第一个[]表示图片数目，第二个[]0表示转化后的tensor，1表示列表，输出后如果是0表示猫，1表示狗，因为cat文件夹在dog上面

batch_size=32
train_loader = data.DataLoader(train_data,batch_size=batch_size,shuffle=True,pin_memory=True)
test_loader = data.DataLoader(test_data,batch_size=batch_size)


class Cnn(nn.Module):
    def __init__(self):
        super(Cnn, self).__init__()
        self.conv1=nn.Conv2d(3,20,5,5)
        self.conv2 = nn.Conv2d(20, 50, 4, 1)
        self.fc1 = nn.Linear(50 * 6 * 6, 200)
        self.fc2 = nn.Linear(200, 2)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        # 20*30*30
        x = F.max_pool2d(x, 2, 2)
        # 20*15*15
        x = F.relu(self.conv2(x))
        # 50*12*12
        x = F.max_pool2d(x, 2, 2)
        # 50*6*6
        x = x.view(-1, 50 * 6 * 6)
        # 压扁成了行向量，(1,50*6*6)
        x = F.relu(self.fc1(x))
        # (1,200)
        x = self.fc2(x)
        # (1,2)
        return F.log_softmax(x, dim=1)
lr=1e-4
device=torch.device("cuda" if torch.cuda.is_available() else "cpu" )
model=Cnn().to(device)
optimizer=optim.Adam(model.parameters(),lr=lr)

def train(model,device,train_loader,optimizer,epoch,losses):
    model.train()
    for idx, (t_data, t_target) in enumerate(train_loader):
        t_data, t_target = t_data.to(device), t_target.to(device)
        pred = model(t_data)  # batch_size*2
        loss = F.nll_loss(pred, t_target)

        # Adam
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if idx % 10 == 0:
            print("epoch:{},iteration:{},loss:{}".format(epoch, idx, loss.item()))
            losses.append(loss.item())


def test(model,device,test_loader):
    model.eval()
    correct = 0  # 预测对了几个。
    with torch.no_grad():
        for idx, (t_data, t_target) in enumerate(test_loader):
            t_data, t_target = t_data.to(device), t_target.to(device)
            pred = model(t_data)  # batch_size*2
            pred_class = pred.argmax(dim=1)  # batch_size*2->batch_size*1
            correct += pred_class.eq(t_target.view_as(pred_class)).sum().item()
    acc = correct / len(test_data)
    print("accuracy:{}".format(acc))


num_epochs=10
losses=[]
from time import *
begin_time=time()
for epoch in range(num_epochs):
    train(model,device,train_loader,optimizer,epoch,losses)
    test(model,device,test_loader)

end_time=time()
torch.save(model, 'model.pth')
model_save_path = 'model.pth'
class_names = ['cat', 'dog']
model = torch.load(model_save_path)
model.eval()
# print(model)

image_PIL = Image.open('21.jpg')
#
image_tensor = transforms(image_PIL)
# 以下语句等效于 image_tensor = torch.unsqueeze(image_tensor, 0)
image_tensor.unsqueeze_(0)
# 没有这句话会报错
image_tensor = image_tensor.to(device)

out = model(image_tensor)
print(torch.argmax(out))

